import os

import torch
import torchvision
import torchvision.transforms as transforms
import numpy
from PIL import Image


class ANN:
    def __init__(self, n, structure):
        self.n = n
        self.structure = structure
        self.a = []
        self.w = []
        #self.b = []
        self.x = []
        self.da = []
        self.dw = []
        #self.db = []
        self.dx = []
        for i in range(n):
            self.a.append(torch.rand((1, structure[i])))
            self.da.append(torch.rand((1, structure[i])))
            if i != 0:
                #self.b.append(torch.rand((1, structure[i])) * 0.001)
                self.w.append(torch.rand((structure[i - 1], structure[i])) * 0.001)
                self.x.append(torch.rand((1, structure[i])))
                #self.db.append(torch.rand((1, structure[i])))
                self.dw.append(torch.rand((structure[i - 1], structure[i])))
                self.dx.append(torch.rand((1, structure[i])))

        self.alpha = 0.0001

    def load(self):
        file = open("test.txt")
        for i in range(self.n - 1):
            for j in range(len(self.w[i])):
                line = file.readline().split()
                for k in range(len(self.w[i][j])):
                    self.w[i][j][k] = float(line[k])
            # line = file.readline().split()
            # for j in range(len(self.b[i][0])):
            #     self.w[i][0][j] = float(line[j])

    def save(self):
        file = open("test.txt", "w")
        for i in range(self.n - 1):
            for j in range(len(self.w[i])):
                for k in range(len(self.w[i][j])):
                    file.write(str(float(self.w[i][j][k])) + " ")
                file.write("\n")
            # for j in range(len(self.b[i][0])):
            #     file.write(str(float(self.b[i][0][j])) + " ")
            # file.write("\n")

    def active(self, input):
        return torch.where(input > 0, input, float(0))

    def dactive(self, input):
        return torch.where(input > 0, float(1), float(0))


    def train(self, dataSet):
        n = 0
        L = 0
        for index, (image, label) in enumerate(dataSet):
            self.a[0] = image.reshape((1, 784))
            for i in range(1, self.n):
                self.x[i - 1] = torch.mm(self.a[i - 1], self.w[i - 1])# + self.b[i - 1]
                self.a[i] = self.active(self.x[i - 1])

            self.da[self.n - 1] = self.a[self.n - 1].clone()
            self.da[self.n - 1][0][label] -= 1
            for i in range(self.n - 1, 0, -1):
                self.dx[i - 1] = self.dactive(self.x[i - 1]) * self.da[i]
                self.dw[i - 1] = self.dx[i - 1].repeat((self.structure[i - 1], 1)) * self.a[i - 1].T.repeat(
                    (1, self.structure[i]))
                #self.db[i - 1] = self.dx[i - 1].clone()
                self.da[i - 1] = torch.mm(self.dx[i - 1], self.w[i - 1].T)

            for i in range(self.n - 1):
                self.w[i] -= self.dw[i] * self.alpha
                #self.b[i] -= self.db[i] * self.alpha

            L += self.da[self.n - 1] @ self.da[self.n - 1].T

            if self.a[self.n - 1][0].argmax() == int(label[0]):
                n += 1

            if int(index) % 1000 == 0:
                print(f"{n}/{int(index) + 1}:{n / (int(index) + 1)}")
                print(f"LOSS:{L / n}")
                print()
                self.save()

class SNN:
    def __init__(self, n, structure):
        self.n = n
        self.structure = structure
        self.x = []
        self.w = []
        self.m = []
        self.v = []
        self.s = []
        self.out = []
        for i in range(n):
            self.x.append(torch.zeros((1, structure[i])))
            self.m.append(torch.zeros((1, structure[i])))
            self.v.append(torch.zeros((1, structure[i])))
            self.s.append(torch.zeros((1, structure[i])))
            if i != 0:
                self.w.append(torch.zeros((self.structure[i - 1], self.structure[i])))

    def init(self):
        for i in range(self.n):
            self.x[i] = torch.zeros((1, self.structure[i]))
            self.m[i] = torch.zeros((1, self.structure[i]))
            self.v[i] = torch.zeros((1, self.structure[i]))
            self.s[i] = torch.zeros((1, self.structure[i]))

    def load(self):
        file = open("ann1.txt")
        for i in range(self.n - 1):
            for j in range(len(self.w[i])):
                line = file.readline().split()
                for k in range(len(self.w[i][j])):
                    self.w[i][j][k] = float(line[k])
            file.readline()

    def run(self, input, time_step):
        r = torch.zeros((1,self.structure[self.n - 1]))
        self.init()
        for i in range(time_step):
            self.x[0] = input.reshape((1, 784))
            for j in range(1, self.n):
                self.m[j] = self.v[j] + torch.mm(self.x[j - 1], self.w[j - 1])
                self.s[j] = torch.where(self.m[j] > 1, float(1), float(0))
                self.x[j] = self.s[j] * 1
                self.v[j] = self.m[j] - self.x[j]
            r += self.x[self.n - 1]
        return r / time_step



if __name__ == "__main__":

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    ann = ANN(3, (784, 100, 10))
    ann.load()
    while True:
        train_dataset = torchvision.datasets.MNIST(root="data", train=True, download=False,
                                                   transform=transforms.ToTensor())
        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)
        ann.train(train_loader)
        ann.save()


    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # train_dataset = torchvision.datasets.MNIST(root="data", train= False, download=False,
    #                                                transform=transforms.ToTensor())
    # train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0)
    #
    # snn = SNN(4, (784, 500, 500, 10))
    # snn.load()
    #
    # n = 0
    # L = 0
    # for i, (image, label) in enumerate(train_loader):
    #     if (int(i) - 1) % 1000 == 0:
    #         print(f"{n}/{int(i)}:{n/int(i)}")
    #     if snn.run(image,16).argmax() == int(label):
    #         n += 1
